use std::collections::HashMap;
use crate::encrypt_tool::protocol::{UserPackage, ConsultMessage, KeyPairDataType, SessionDataType, RatchetRootKey, RatchetChainKey};

mod protocol;

use crypto;
use rand::Rng;
use crypto::ed25519::keypair;
use std::iter::repeat;

/// 加解密工具
/// encrypt_msg 将原文 -> 密文
/// decrypt_msg 将密文 -> 原文
trait SecretTool
{
    // 加密消息
    fn encrypt_msg(&self, fid: i32, original_text: &str) -> (bool, String);
    // 解密消息
    fn decrypt_msg(&self, fid: i32, cipher_text: &str) -> (bool, String);
}

/// 协商工具
/// build_consult_message 发起方 构建协商消息
/// analyze_consult_message 接受方解析协商信息
pub trait ConsultTool {
    fn build_consult_message(&mut self, user: UserPackage) -> ConsultMessage;
    fn analyze_consult_message(&mut self, consult: ConsultMessage);
}

pub trait KeyHelper {
    /// 只生成 私钥
    fn gen_pri_key(&self) -> Vec<u8>;
    /// 只生成 公钥
    fn gen_pub_key(&self, pri_key: &Vec<u8>) -> Vec<u8>;
    /// 生成u一对公私钥对
    fn gen_key_pair(&self) -> KeyPairDataType;
}

pub trait HKDFHelper{
    ///
    /// 生成下一圈棘轮
    /// salt 盐
    /// ikm 输入材料
    /// info 内容
    /// okm 输出
    fn gen_next_key(&self , salt: &[u8], ikm: &[u8],  info: &[u8] ,okm: &mut [u8]);
}

/// 加密工具 一个用户一套
pub struct EncrtptTool
{
    // 储存 tcpid 和 本地缓存 的用户信息
    ///< uid , 棘轮
    friend_map: HashMap<i32, SessionDataType>,
    opk_map: HashMap<u32, KeyPairDataType>,
    my_uid: i32,
    my_ik: Option<KeyPairDataType>,
    my_spk: Option<KeyPairDataType>,
    my_sign: Option<KeyPairDataType>,
}


impl EncrtptTool {
    pub fn new() -> EncrtptTool {
        EncrtptTool
        {
            friend_map: Default::default(),
            opk_map: Default::default(),
            my_uid: 0,
            my_ik: None,
            my_spk: None,
            my_sign: None,
        }
    }

    pub fn setUid(&mut self, uid: i32) {
        self.my_uid = uid;
    }
    fn setIk(&mut self, keypair: KeyPairDataType)
    {
        self.my_ik = Some(keypair);
    }

    fn setSPK(&mut self, keypair: KeyPairDataType)
    {
        self.my_spk = Some(keypair);
    }
    // 生成签名
    fn genSign(&mut self) {
        // TODO wujehy 暂时 你不做签名处理
    }


    pub fn init_user(&mut self, ikpair: KeyPairDataType, spkpair: KeyPairDataType) {
        self.setIk(ikpair);
        self.setSPK(spkpair);
        self.genSign();
    }
    fn jenkins_hash(&self, data: &Vec<u8>) -> u32 {
        // 控制 数值在 32 位
        let mut result = 0u128;
        for v in data.iter() {
            result = result & 0xffff_ffff_ffff_ffff;
            result += v.clone() as u128;
            let temp = (result & 0xffff_ffff_ffff_ffff) << 10;
            result += temp;
            result ^= ((result & 0xffff_ffff_ffff_ffff) >> 6);
        }
        result += ((result & 0xffff_ffff_ffff_ffff) << 3);
        result ^= ((result & 0xffff_ffff_ffff_ffff) >> 11);
        result += ((result & 0xffff_ffff_ffff_ffff) << 15);
        result as u32
    }
    pub fn gen_opk(&mut self, size: i32) {
        for _ in 0..size {
            let tmp = self.gen_key_pair();
            let getKid = self.jenkins_hash(&tmp.pub_key);
            println!("add opk id {} ", getKid);
            self.opk_map.insert(getKid, tmp);
        }
    }


    //
    pub fn pop_user_package(&mut self) -> UserPackage {
        let mut getopk_pub = vec![];

        for (opkid, opk) in self.opk_map.iter() {
            getopk_pub = opk.pub_key.clone();
        }

        // 如果没有了 则临时生成
        if getopk_pub.len() == 0 {
            let newopk = self.gen_key_pair();
            let getKid = self.jenkins_hash(&newopk.pub_key);
            println!("add opk id {} ", getKid);
            getopk_pub = newopk.pub_key.clone();
            self.opk_map.insert(getKid, newopk);
        }

        UserPackage {
            uid: self.my_uid,
            other_ik_pub: self.my_ik.as_ref().unwrap().pub_key.clone(),
            other_spk_pub: self.my_spk.as_ref().unwrap().pub_key.clone(),
            other_opk_pub: getopk_pub,
        }
    }
}

impl HKDFHelper for EncrtptTool
{
    fn gen_next_key(&self ,salt: &[u8], ikm: &[u8], info: &[u8], okm: &mut [u8]) {
        let mut prk: Vec<u8> = repeat(0).take(32).collect();
        let digest =crypto::sha2::Sha256::new();
        crypto::hkdf::hkdf_extract(digest, &salt[..], &ikm[..], &mut prk);
        crypto::hkdf::hkdf_expand(digest, &prk[..], &info[..], okm);
    }
}

impl ConsultTool for EncrtptTool
{
    fn build_consult_message(&mut self, user: UserPackage) -> ConsultMessage {
        // TODO 协商
        // 生成一个临时密钥
        let epk = self.gen_key_pair();

        // 准备材料
        let epk_pir = epk.pri_key;
        let myik_pri = self.my_ik.as_ref().unwrap().pri_key.clone();
        let other_ik = user.other_ik_pub;
        let other_spk = user.other_spk_pub;
        let other_opk = user.other_opk_pub;

        let opkid = self.jenkins_hash(&other_opk);
        let dh1 = crypto::curve25519::curve25519(myik_pri.as_slice(), other_spk.as_slice());
        println!("k1 = {:?} , k2 ={:?}", myik_pri.as_slice(), other_spk.as_slice());
        println!("dh 1 : {:?} ", dh1);
        let dh2 = crypto::curve25519::curve25519(epk_pir.as_slice(), other_ik.as_slice());
        println!("dh 2 : {:?} ", dh2);
        let dh3 = crypto::curve25519::curve25519(epk_pir.as_slice(), other_spk.as_slice());
        println!("dh 3 : {:?} ", dh3);
        let dh4 = crypto::curve25519::curve25519(epk_pir.as_slice(), other_opk.as_slice());
        println!("dh 4 : {:?} ", dh4);

        // 使用 4个DH 生成摘要 这里使用 kdf 生成 64 位 摘要

        let mut master_root_key =  vec![];
        for &v in dh1.iter() {
            master_root_key.push(v);
        }
        for &v in dh2.iter() {
            master_root_key.push(v);
        }
        for &v in dh3.iter() {
            master_root_key.push(v);
        }
        for &v in dh4.iter() {
            master_root_key.push(v);
        }

        /// NOTE wujehy 初始盐 都是32 位 的0 , 之后是epk
        /// ikm 是上一次导出的前32位
        /// info 可以是 内容 master_root_key

        println!("master_root_key {:?}",master_root_key);
        let salt  = vec![];
        let info=master_root_key  ;
        let ikm = vec![] ;
        let mut okm: Vec<u8> = repeat(0).take(64).collect();

        self.gen_next_key(&salt[..], &ikm[..],&info[..],&mut okm );

        let new_session = SessionDataType{
            uid: user.uid,
            rootkey: RatchetRootKey {
                key:  okm
            },
            chainkey_sender: RatchetChainKey { key: vec![], index: 0 },
            chainkey_receiver: RatchetChainKey {key: vec![], index: 0},
            message_history: Default::default()
        } ;

        // 将会话信息 添加到 好友列表
        self.friend_map.insert(user.uid , new_session);

        /// NOTE wujehy 传输的时候只能使用 公钥传输
        ConsultMessage {
            uid: self.my_uid,
            other_ik_pub: self.my_ik.as_ref().unwrap().pub_key.clone(),
             ///< 我的 ik 公钥
            other_epk_pub: epk.pub_key,
             ///< 临时生成的 epk 公钥
            opkid: opkid,
            body: vec![],
        }
    }

    fn analyze_consult_message(&mut self, consult: ConsultMessage) {
        /// 接受方解析协商信息
        ///
        /// 数据准备
        let my_ik_pir = self.my_ik.as_ref().unwrap().pri_key.clone();
        let my_spk_pir = self.my_spk.as_ref().unwrap().pri_key.clone();
        // TODO opk 实际情况用完就删 都则 这里测试 , 不需要
        // let mut my_opk_pair  = self.opk_map.remove(&consult.opkid).unwrap(); /// 用完就删
        let mut my_opk_pair = self.opk_map.get(&consult.opkid).unwrap();
        let my_opk_pir = my_opk_pair.pri_key.clone();

        let other_ik_pub = consult.other_ik_pub;
        let other_epk_pub = consult.other_epk_pub;

        // 计算 dh

        // 这里输出的 dh 与发送方的一致
        println!("k1 = {:?} , k2 ={:?}", my_spk_pir.as_slice(), other_ik_pub.as_slice());
        let dh1 = crypto::curve25519::curve25519(my_spk_pir.as_slice(), other_ik_pub.as_slice());

        println!("dh 1 : {:?} ", dh1);
        let dh2 = crypto::curve25519::curve25519(my_ik_pir.as_slice(), other_epk_pub.as_slice());
        println!("dh 2 : {:?} ", dh2);
        let dh3 = crypto::curve25519::curve25519(my_spk_pir.as_slice(), other_epk_pub.as_slice());
        println!("dh 3 : {:?} ", dh3);
        let dh4 = crypto::curve25519::curve25519(my_opk_pir.as_slice(), other_epk_pub.as_slice());
        println!("dh 4 : {:?} ", dh4);
        // 通过id  查询是否有这个opk

        let mut master_root_key =  vec![];
        for &v in dh1.iter() {
            master_root_key.push(v);
        }
        for &v in dh2.iter() {
            master_root_key.push(v);
        }
        for &v in dh3.iter() {
            master_root_key.push(v);
        }
        for &v in dh4.iter() {
            master_root_key.push(v);
        }

        /// NOTE wujehy 初始盐 都是32 位 的0 , 之后是epk
        /// ikm 是上一次导出的前32位
        /// info 可以是 内容 master_root_key

        println!("master_root_key {:?}",master_root_key);
        let salt  = vec![];
        let info=master_root_key  ;
        let ikm = vec![] ;
        let mut okm: Vec<u8> = repeat(0).take(64).collect();

        self.gen_next_key(&salt[..], &ikm[..],&info[..],&mut okm );

        let new_session = SessionDataType{
            uid: consult.uid,
            rootkey: RatchetRootKey {
                key:  okm
            },
            chainkey_sender: RatchetChainKey { key: vec![], index: 0 },
            chainkey_receiver: RatchetChainKey {key: vec![], index: 0},
            message_history: Default::default()
        } ;

        self.friend_map.insert(consult.uid , new_session);
    }
}

impl KeyHelper for EncrtptTool {
    fn gen_pri_key(&self) -> Vec<u8> {
        // 使用随机数生成 32 位长度的key
        let mut rng = rand::thread_rng();
        let mut prikey = Vec::new();
        for _ in 1..32 {
            let data = rng.gen::<u8>();
            // println!("nub : {} " , data);
            prikey.push(data);
        }
        // println!("key : {:?}" , prikey);
        prikey
    }

    fn gen_pub_key(&self, pri_key: &Vec<u8>) -> Vec<u8> {
        let pubkey = crypto::curve25519::curve25519_base(pri_key);
        pubkey.to_vec()
    }

    fn gen_key_pair(&self) -> KeyPairDataType
    {
        let prikey = self.gen_pri_key();
        let pubkey = self.gen_pub_key(prikey.clone().as_ref());
        // println!("prikey : {:?} , pubkey :{:?}" , prikey , pubkey);
        KeyPairDataType::new(prikey, pubkey)
    }
}

impl SecretTool for EncrtptTool{
    fn encrypt_msg(&self, fid: i32, original_text: &str) -> (bool, String) {
        /// TODO wujehy 发送棘轮转动
        /// 拿到这个好友的 信息
        /// NOTE 暂时不做异常处理
        let finfo = self.friend_map.get(&fid).unwrap();

        let noindex = 0;
        if (finfo.chainkey_sender.index == noindex) {
            /// 使用 root key 开始


            (true , String::from("ok"))
        }else {
            /// 使用 sendkey 最后一个 计算
            (false , String::from("error"))

        }

    }

    fn decrypt_msg(&self, fid: i32, cipher_text: &str) -> (bool, String) {
        // TODO wujehy 接受棘轮转动


        unimplemented!()
    }
}